#include "Trajectory.h"

PathPoint::PathPoint(dReal timeStep, const dReal* pos, const dReal* rot, const dReal* mousePos)
{
	this->timeStep = timeStep;
	for(int i = 0; i < 3; i++)
		this->pos[i] = pos[i];
	for(int i = 0; i < 4; i++)
		this->rot[i] = rot[i];
	for(int i = 0; i < 3; i++)
		this->mousePos[i] = mousePos[i];
}

void Trajectory::record()
{
	pathPoints.clear();
	totalInfos.clear();
}

void Trajectory::saveE(QString filename)
{
	QFile fileE(filename);
	if(!fileE.open(QIODevice::WriteOnly | QIODevice::Text))
		qDebug() << "cannot write file E !";

	QTextStream outE(&fileE);

	for(QVector<PathPoint>::iterator iter = pathPoints.begin(); iter != pathPoints.end(); iter++)
	{
		dReal timeStep = iter->timeStep;
		dVector3 pos = {iter->pos[0], iter->pos[1], iter->pos[2]};
		dQuaternion rot = {iter->rot[0], iter->rot[1], iter->rot[2], iter->rot[3]};
		dVector3 mousePos = {iter->mousePos[0], iter->mousePos[1], iter->mousePos[2]};
		outE << timeStep << "," << pos[0] << "," << pos[1] << "," << pos[2] << "," << rot[0] << "," << rot[1] << "," << rot[2] << "," << rot[3] << "," << mousePos[0] << "," << mousePos[1] << "," << mousePos[2] << "\n";
	}

	fileE.close();
}

void Trajectory::addMotorNoise()
{
	Eigen::MatrixXd dataMatrix(totalInfos.size(), 17);
	int count = 0;
	for(std::vector<TotalInfo, Eigen::aligned_allocator<TotalInfo>> ::iterator iter = totalInfos.begin(); iter != totalInfos.end(); iter++)
	{
		dReal timeStep = iter->timeStep;
		dReal weight = iter->weight;
		dReal mu = iter->mu;
		Vector3d force(iter->force);
		Vector3d torque(iter->torque);
		Vector3d pos(iter->pos);
		Quaternion<double> rot(iter->rot);
		int status = iter->status;
		Vector3d conForce(iter->conForce);

		for(int i = 0; i < 3; i++)
			dataMatrix(count, i) = force[i];
		for(int i = 0; i < 3; i++)
			dataMatrix(count, i+3) = torque[i];
		for(int i = 0; i < 3; i++)
			dataMatrix(count, i+6) = pos[i];
		dataMatrix(count, 9) = rot.w();
		dataMatrix(count, 10) = rot.x();
		dataMatrix(count, 11) = rot.y();
		dataMatrix(count, 12) = rot.z();
		dataMatrix(count, 13) = status;
		for(int i = 0; i < 3; i++)
			dataMatrix(count, i+14) = conForce[i];
		count++;
	}

	VectorXd maxCoeff = dataMatrix.colwise().maxCoeff();
	VectorXd minCoeff = dataMatrix.colwise().minCoeff();

	for(std::vector<TotalInfo, Eigen::aligned_allocator<TotalInfo>> ::iterator iter = totalInfos.begin(); iter != totalInfos.end(); iter++)
	{
		dReal timeStep = iter->timeStep;
		dReal weight = iter->weight;
		dReal mu = iter->mu;
		Vector3d force(iter->force);
		Vector3d torque(iter->torque);
		Vector3d pos(iter->pos);
		Quaternion<double> rot(iter->rot);
		int status = iter->status;
		Vector3d conForce(iter->conForce);

		if(Macro::s_motorNoise && Macro::s_replay)
		{
			std::random_device rd;
			std::default_random_engine generator(rd());

			//add noise to force
			for(int i = 0; i < 3; i++)
			{
				double upper = maxCoeff[i];
				double lower = minCoeff[i];
				std::normal_distribution<double> distribution(0.0, 0.05 * (upper - lower) + 1e-4);

				while(1)
				{
					double gaussianNoise = distribution(generator);
					if(gaussianNoise < 0.005 && gaussianNoise > -0.005)
					{
						force[i] += gaussianNoise;
						break;
					}
				}
			}

			//add noise to torque
			for(int i = 0; i < 3; i++)
			{
				double upper = maxCoeff[i+3];
				double lower = minCoeff[i+3];
				std::normal_distribution<double> distribution(0.0, 0.05 * (upper - lower) + 1e-4);

				while(1)
				{
					double gaussianNoise = distribution(generator);
					if(gaussianNoise < 0.005 && gaussianNoise > -0.005)
					{
						torque[i] += gaussianNoise;
						break;
					}
				}
			}

			//add noise to pos
			for(int i = 0; i < 3; i++)
			{
				double upper = maxCoeff[i+6];
				double lower = minCoeff[i+6];
				std::normal_distribution<double> distribution(0.0, 0.05 * (upper - lower) + 1e-4);
				while(1)
				{					
					double gaussianNoise = distribution(generator);
					if(gaussianNoise < 0.005 && gaussianNoise > -0.005)
					{
						pos[i] += gaussianNoise;
						break;
					}
				}
			}

			//add noise to quat
			{
				double upper = maxCoeff[9];
				double lower = minCoeff[9];
				std::normal_distribution<double> distribution(0.0, 0.05 * (upper - lower) + 1e-4);
				while(1)
				{
					double gaussianNoise = distribution(generator);
					if(gaussianNoise < 0.005 && gaussianNoise > -0.005)
					{
						rot.w() += gaussianNoise;
						break;
					}
				}
			}
			{
				double upper = maxCoeff[10];
				double lower = minCoeff[10];
				std::normal_distribution<double> distribution(0.0, 0.05 * (upper - lower) + 1e-4);
				while(1)
				{
					double gaussianNoise = distribution(generator);
					if(gaussianNoise < 0.005 && gaussianNoise > -0.005)
					{
						rot.x() += gaussianNoise;
						break;
					}
				}
			}
			{
				double upper = maxCoeff[11];
				double lower = minCoeff[11];
				std::normal_distribution<double> distribution(0.0, 0.05 * (upper - lower) + 1e-4);
				while(1)
				{
					double upper = maxCoeff[11];
					double lower = minCoeff[11];
					std::default_random_engine generator;
					std::normal_distribution<double> distribution(0.0, 0.05 * (upper - lower) + 1e-4);

					double gaussianNoise = distribution(generator);
					if(gaussianNoise < 0.005 && gaussianNoise > -0.005)
					{
						rot.y() += gaussianNoise;
						break;
					}
				}
			}
			{
				double upper = maxCoeff[12];
				double lower = minCoeff[12];
				std::normal_distribution<double> distribution(0.0, 0.05 * (upper - lower) + 1e-4);
				while(1)
				{
					double gaussianNoise = distribution(generator);
					if(gaussianNoise < 0.005 && gaussianNoise > -0.005)
					{
						rot.z() += gaussianNoise;
						break;
					}
				}
			}

			//add noise to contact force
			if(status == 1)
			{
				for(int i = 0; i < 3; i++)
				{
					double upper = maxCoeff[i+14];
					double lower = minCoeff[i+14];
					std::normal_distribution<double> distribution(0.0, 0.05 * (upper - lower) + 1e-4);

					while(1)
					{
						double gaussianNoise = distribution(generator);
						if(gaussianNoise < 0.005 && gaussianNoise > -0.005)
						{
							conForce[i] += gaussianNoise;
							break;
						}
					}
				}
			}
		}

		iter->force = force;
		iter->torque = torque;
		iter->pos = pos;
		iter->rot = rot;
		iter->conForce = conForce;
	}
}

void Trajectory::saveT(QString filename)
{
	QFile fileT(filename);
	if(!fileT.open(QIODevice::WriteOnly | QIODevice::Text))
		qDebug() << "cannot write file W !";

	QTextStream outW(&fileT);

	addMotorNoise();

	for(std::vector<TotalInfo, Eigen::aligned_allocator<TotalInfo>> ::iterator iter = totalInfos.begin(); iter != totalInfos.end(); iter++)
	{
		dReal timeStep = iter->timeStep;
		dReal weight = iter->weight;
		dReal mu = iter->mu;
		Vector3d force(iter->force);
		Vector3d torque(iter->torque);
		Vector3d pos(iter->pos);
		Quaternion<double> rot(iter->rot);
		int status = iter->status;
		Vector3d conForce(iter->conForce);
		outW << timeStep << "," << weight << "," << mu << ","
			<< force[0] << "," << force[1] << ","  << force[2] << ","  
			<< torque[0] << "," << torque[1] << ","  << torque[2] << ","  
			<< pos[0] << ","  << pos[1] << ","  << pos[2] << ","  
			<< rot.w() << ","  << rot.x() << ","  << rot.y() << ","  << rot.z() << ","
			<< status << "," << conForce[0] << "," << conForce[1] << ","  << conForce[2]
			<< "\n";
	}

	fileT.close();
}

void Trajectory::load(QString filename)
{
	pathPoints.clear();

	QFile file(filename);
	if (!file.open(QIODevice::ReadOnly | QIODevice::Text)){
		qDebug()<< "cannot read file!";     
	}

	Csv csv(&file);
	while(1)
	{
		QStringList line = csv.parseLine();
		if(line.size() == 1)
			break;
		else
		{
			dReal timeStep = line.at(0).toDouble();
			dVector3 pos = {line.at(1).toDouble(), line.at(2).toDouble(), line.at(3).toDouble()};
			dQuaternion rot = {line.at(4).toDouble(), line.at(5).toDouble(), line.at(6).toDouble(), line.at(7).toDouble()};
			dVector3 mousePos = {line.at(8).toDouble(), line.at(9).toDouble(), line.at(10).toDouble()};
			PathPoint pathPoint(timeStep, pos, rot, mousePos);
			pathPoints.push_back(pathPoint);
		}
	}

	file.close();
}

void Trajectory::pause()
{

}

void Trajectory::reset()
{
	pathPoints.clear();
	totalInfos.clear();
}